import numpy as np
import operator

class KNN(object):

    def __init__(self, k=3):

        self.k = k
        self.x = None
        self.y = None

    def fit(self, x, y):
        self.x = x
        self.y = y

    def square_distance(self, v1, v2):
        return np.sum(np.square(v1-v2))

    def vote(self, ys):

        # ys_unique = np.unique(ys)

        vote_dict = {}

        for y in ys:
            if y not in vote_dict.keys():
                vote_dict[y] = 1
            else:
                vote_dict[y] += 1
            sorted_vote_dict = sorted(vote_dict.items(),
                                      key=operator.itemgetter(1),
                                      reverse=True)
            return sorted_vote_dict[0][0]

    def predict(self, x):
        y_pred = []

        for i in range(len(x)):
            dist_array = [self.square_distance(x[i], self.x[j]) for j in range(len(self.x))]
            sorted_index = np.argsort(dist_array)
            top_k_index = sorted_index[:self.k]
            y_pred.append(self.vote(ys=self.y[top_k_index]))
        return np.array(y_pred)

    def score(self, y_true=None, y_pred=None):

        if y_true is None or y_pred is None:
            y_pred = self.predict(self.x)
            y_true = self.y

        score = 0.0

        for i in range(len(y_true)):
            if y_true[i].all() == y_pred[i]:
                score += 1

        score /= len(y_true)

        return score


